{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ab03a19a",
   "metadata": {},
   "source": [
    "# 注册算子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5cf71110",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "\n",
    "def register(op_name, describe=\"\"):\n",
    "    \"\"\"Get the Op for a given name.\n",
    "    when the op_name is not registered, create a new empty op with the given name.\n",
    "    when the op_name has been registered, abort with an error message.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    op_name : str\n",
    "        The operator name\n",
    "\n",
    "    describe : Optional[str]\n",
    "        The operator description\n",
    "    \"\"\"\n",
    "\n",
    "    tvm.ir._ffi_api.RegisterOp(op_name, describe)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d3c57f5",
   "metadata": {},
   "source": [
    "## 1. 获取或创建算子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8878cce2",
   "metadata": {},
   "outputs": [],
   "source": [
    "op_name = \"my.operator\"\n",
    "try:\n",
    "    op = tvm.ir.Op.get(op_name)\n",
    "except:\n",
    "    op = register(op_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02f88ca2",
   "metadata": {},
   "source": [
    "## 2. 设置算子属性"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cea4e8f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "tvm.ir.Op.get(op_name).set_num_inputs(1)\n",
    "tvm.ir.Op.get(op_name).add_argument(\"data\", \"Tensor\", \"输入数据\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e945a752",
   "metadata": {},
   "source": [
    "## 3. 定义类型关系"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b92a99c6",
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'Op' object has no attribute 'add_type_rel'",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mAttributeError\u001b[39m                            Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[4]\u001b[39m\u001b[32m, line 8\u001b[39m\n\u001b[32m      6\u001b[39m     reporter.assoc(types[\u001b[32m1\u001b[39m], data)\n\u001b[32m      7\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[32m----> \u001b[39m\u001b[32m8\u001b[39m \u001b[43mtvm\u001b[49m\u001b[43m.\u001b[49m\u001b[43mir\u001b[49m\u001b[43m.\u001b[49m\u001b[43mOp\u001b[49m\u001b[43m.\u001b[49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[43mop_name\u001b[49m\u001b[43m)\u001b[49m\u001b[43m.\u001b[49m\u001b[43madd_type_rel\u001b[49m(\u001b[33m\"\u001b[39m\u001b[33mMyOpTypeRel\u001b[39m\u001b[33m\"\u001b[39m, my_op_type_rel)\n",
      "\u001b[31mAttributeError\u001b[39m: 'Op' object has no attribute 'add_type_rel'"
     ]
    }
   ],
   "source": [
    "def my_op_type_rel(types, num_inputs, attrs, reporter):\n",
    "    # 验证输入类型并关联输出类型\n",
    "    data = types[0].as_tensor_type()\n",
    "    if data is None:\n",
    "        return False\n",
    "    reporter.assoc(types[1], data)\n",
    "    return True\n",
    "tvm.ir.Op.get(op_name).add_type_rel(\"MyOpTypeRel\", my_op_type_rel)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ff01a71",
   "metadata": {},
   "source": [
    "## 4. 设置计算函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c4a4ae4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def my_op_compute(attrs, inputs, out_type):\n",
    "    # 从TVMScript模块获取计算函数\n",
    "    mod = my_op_module()\n",
    "    # 创建外部函数调用\n",
    "    out = te.extern(\n",
    "        [inputs[0].shape],\n",
    "        [inputs[0]],\n",
    "        lambda ins, outs: mod[\"my_op_func\"](ins[0], outs[0]),\n",
    "        name=\"my_op\",\n",
    "        dtype=inputs[0].dtype\n",
    "    )\n",
    "    return [out]\n",
    "tvm.ir.Op.get(op_name).set_attr(\"FTVMCompute\", my_op_compute)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
